#include "header.h"

static fftw_complex *MiRk;
static fftw_plan plan_f1, plan_b1;
static fftw_plan plan_fr1, plan_br1;
static double *k, *u, overN, q, *mk, dtau, gama, sp, num, den, err, err2, l2, tol, hmax, ml0, ml, htt;
static double coeff, aa, bb, NV, *ra[2], *rt1, *rt2, *rhu[1], *rhk[2], *redk[2], *reDk, *rekD, *reD, *rMiRk;
static double ip, *reDk, *reRk, *reDyk;
static int N, n2; 
static work_ptr tmp;

void err_msg(char* str)
{
  printf("%s\n",str);
  exit(1);
}

double signum(double x) {
  if (x > 0) return  1.;
  if (x < 0) return -1.;
  return 0.;
}

double signum2(double x) {
  if (x > 0) return  1.;
  else return -1.;
}

void init(int *argc, char **argv, params_ptr in, work_ptr wrk)
{
  char str[480], line[480], value[480], param[480];
  sprintf(str, "Usage:\n\t%s input", argv[0]);
  if (*argc != 2) err_msg(str);

  sprintf(str, "%s.cfg", argv[1]);
  FILE *fh = fopen(str,"r"); 
  FILE *fhlog = fopen("run.log","w"); 
  if (fh == NULL) err_msg("Cannot open file");
  while (fgets(line, 480, fh)!=NULL) {
    sscanf(line, "%s\t%s", param, value);
    if (strcmp(param,"#runname=") == 0) sprintf(in->runname,"%s", value);
    if (strcmp(param,"#resname=") == 0) sprintf(in->resname,"%s", value);
    if (strcmp(param,"#npoints=") == 0) in->n = atoi(value);
    if (strcmp(param,"#gravity=") == 0) in->g = atof(value);
    if (strcmp(param,"#velocit=") == 0) in->c = atof(value);
    if (strcmp(param,"#delta_t=") == 0) in->dtau = atof(value);
    if (strcmp(param,"#toleran=") == 0) in->tol = atof(value);
    if (strcmp(param,"#perturb=") == 0) in->alfa = atof(value);
    if (strcmp(param,"#io_time=") == 0) in->t_skip = atoi(value);
    if (strcmp(param,"#hmax_tt=") == 0) htt = atof(value);
  }
  fprintf(fhlog, "runname = %s\n", in->runname );
  fprintf(fhlog, "resname = %s\n", in->resname );
  fprintf(fhlog, "npoints = %d\n", in->n );
  fprintf(fhlog, "gravity = %f\n", in->g );
  fprintf(fhlog, "velocit = %f\n", in->c );
  fprintf(fhlog, "toleran = %f\n", in->tol );
  fprintf(fhlog, "delta_t = %f\n", in->dtau );
  fprintf(fhlog, "perturb = %f\n", in->alfa );
  fprintf(fhlog, "hmax_tt = %f\n", htt );
  fprintf(fhlog, "io_time = %d\n", in->t_skip );
  fclose(fh); fclose(fhlog);
  N = in->n; n2 = N/2 + 1;

  ra[0] = fftw_malloc(sizeof(double)*n2);
  ra[1] = fftw_malloc(sizeof(double)*n2);
  wrk->YY  = fftw_malloc(sizeof(double)*n2);
  wrk->YYk = fftw_malloc(sizeof(double)*n2);

  rhu[0] = fftw_malloc(sizeof(double)*n2);
  rhk[0] = fftw_malloc(sizeof(double)*n2);
  rhk[1] = fftw_malloc(sizeof(double)*n2);

  redk[0] = fftw_malloc(sizeof(double)*n2);
  redk[1] = fftw_malloc(sizeof(double)*n2);

  reDk = fftw_malloc(sizeof(double)*n2);
  rekD = fftw_malloc(sizeof(double)*n2);
  reD  = fftw_malloc(sizeof(double)*n2);

  reDyk = fftw_malloc(sizeof(double)*n2);
  reRk  = fftw_malloc(sizeof(double)*n2);
  reDk  = fftw_malloc(sizeof(double)*n2);
  rMiRk = fftw_malloc(sizeof(double)*n2);

  u  = fftw_malloc(sizeof(double)*(in->n));
  q  = (in->g)/pow(in->c,2);

  if (fftw_import_wisdom_from_filename("r2r.fftw") == 0) {
    printf("Creating plans\n");
    fftw_set_timelimit(60.);
    plan_fr1 = fftw_plan_r2r_1d(n2, ra[0], ra[1], FFTW_REDFT00, fmode);
    plan_br1 = fftw_plan_r2r_1d(n2, ra[1], ra[0], FFTW_REDFT00, fmode);
    fftw_export_wisdom_to_filename("r2r.fftw");
    err_msg("Plans created. Rerun the computation");
  } else {
    fftw_set_timelimit(60.);
    plan_fr1 = fftw_plan_r2r_1d(n2, ra[0], ra[1], FFTW_REDFT00, fmode);
    plan_br1 = fftw_plan_r2r_1d(n2, ra[1], ra[0], FFTW_REDFT00, fmode);
  }
  for (int j = 0; j < in->n; j ++) {
    u[j] = 2*pi*j/(in->n);

    //wrk->Y[j] = 0.2*exp(-2*pow(u[j]-pi,2)); // -cos(3.*u[j])
  }
  fh = fopen(in->resname,"r");
  int jj = 0;
  if (fh != NULL) {
    char v1[80],v2[80],v3[80],v4[80], *dummy;
    while (fgets(line, 480, fh)!=NULL) {
      if (jj == 0) {
        dummy = fgets(line, 160, fh);
        dummy = fgets(line, 160, fh);
      }
      sscanf(line, "%s\t%s\t%s\t%s\n",  v1, v2, v3, v4);
      if (jj == n2) break;
      wrk->YY[jj] = atof(v4);
      jj++;
      
    }
    printf("Lines read = %d of %d expected\n", jj, n2);
    if ( jj != n2) err_msg("I/O error\n");
    fclose(fh);
  }
  

  overN = 1./N; 
  in->w_ptr = wrk; dtau = in->dtau;
  gama = 1. + 1./dtau; tol = 5.e-10;
  fft_even(wrk->YY,wrk->YYk);
  for(int j = 0; j < n2; j++) wrk->YYk[j] = (wrk->YYk[j])*overN*cexp((in->alfa)*j); 

  /*if (jj == 0) {
    wrk->Yk[0] = 0;
    for (int j = 1; j < N; j++) wrk->Yk[0] += -cabs(k[j])*conj(wrk->Yk[j])*(wrk->Yk[j]); 
  }*/
}

void khat_even(double *in, double *out) {
  memcpy(ra[0], in, sizeof(double)*n2);
  writer_out(ra[0], "orig_r2r.dat");
  fftw_execute(plan_fr1);
  writer_out(ra[1], "fft_r2r.dat");
  for (int j = 0; j < n2 ; j++) ra[1][j] = j*overN*ra[1][j];
  fftw_execute(plan_br1);
  memcpy(out, ra[0], sizeof(double)*n2);
}

void fft_even(double *in, double *out) {
  memcpy(ra[0], in, n2*sizeof(double));
  fftw_execute(plan_fr1);
  memcpy(out, ra[1], n2*sizeof(double));
}

void ifft_even(double *in, double *out) {
  memcpy(ra[1], in, n2*sizeof(double));
  fftw_execute(plan_br1);
  memcpy(out, ra[0], n2*sizeof(double));
}

double inner_product_real(double *in1, double *in2) {
  ip = 0.;
  for (int j = 0; j < n2; j++) {
    ip += 2.*in1[j]*in2[j];
  }
  ip = ip - (in1[0]*in2[0] + in1[n2-1]*in2[n2-1]);
  return ip;
}

void prepare_rhs_even(work_ptr wrk) {
  ifft_even(wrk->YYk,wrk->YY); 
  for (int j = 0; j < n2; j++) {
    rhk[0][j] = (wrk->YYk[j])*j;
    rhk[1][j] = (wrk->YY[j])*(wrk->YY[j]);
  }
  fft_even(rhk[1],rhk[1]);  
  ifft_even(rhk[0],rhu[0]);
  for (int j = 0; j < n2; j++) {
    rhk[1][j] = j*rhk[1][j]*overN;   
    rhk[0][j] = rhu[0][j]*(wrk->YY[j])*overN;
  }
  fft_even(rhk[0],rhk[0]); 
  for (int j = 0; j < n2; j++) {
    rhk[0][j] = -(j - q)*(wrk->YYk[j]) + q*(0.5*rhk[1][j] + rhk[0][j]); 
  }
}

void prepare_linearization_even(work_ptr wrk) {
  ifft_even(reDk, reD);
  for (int j = 0; j < n2; j++) {
    redk[0][j] = (wrk->YY[j])*reD[j];   
    redk[1][j] = j*reDk[j];    
  }
  fft_even(redk[0], redk[0]);
  ifft_even(redk[1], rekD);
  for (int j = 0; j < n2; j++) {
    redk[0][j] = j*redk[0][j]*overN;
    redk[1][j] = (wrk->YY[j])*rekD[j] + rhu[0][j]*reD[j];
  }
  fft_even(redk[1], redk[1]);
  for (int j = 0; j < n2; j++) {
    redk[1][j] = redk[1][j]*overN;
    redk[0][j] = -(j - q)*reDk[j] + q*(redk[0][j] + redk[1][j]);
  }
}

void iterate_cg(work_ptr wrk) {
  int counter, counter2;
  double cg_error = 1.;
  double err = 1.;
  double eps = 0.02;

  counter2 = 1; NV = 1.; counter = 0;
  while (NV > tol) {
    prepare_rhs_even(wrk);
    writer_out(wrk->YY,"stokes_phys.dat");
    writer_out(wrk->YY,"restart.dat");
    for (int j = 0; j < n2; j++) {
      reDyk[j] = 0;
      reRk[j] = -1.*rhk[0][j];
      reDk[j] = reRk[j]/(j - q);
    }
     //sqrt(inner_product_real(reRk,reRk));
    NV = sqrt(inner_product_real(rhk[0],rhk[0]) );  
    if ((counter2%1) == 0 ) {
        hmax = ((wrk->YY[n2-1])-(wrk->YY[0]))/(2.*pi);
	printf("Nevyazka = %e Newton cycle %1d\n", NV, counter2);
    }
    counter = 0; cg_error = 1.;

    while (cg_error > eps*NV ) {       // start CG cycle 
      prepare_linearization_even(wrk);
      for (int j = 0; j < n2; j++) rMiRk[j] = reRk[j]/(j - q);
      num = inner_product_real(reRk,rMiRk);
      den = inner_product_real(reDk,redk[0]);
      aa = num/den;
      for (int j = 0; j < n2; j++) {
        reDyk[j]  = reDyk[j] + aa*reDk[j];
        rhk[1][j] = reRk[j];                       
        reRk[j]   = reRk[j] - aa*redk[0][j];
        rhk[0][j] = reRk[j]/(j - q);
      }
      num = inner_product_real(reRk, rhk[0]);
      den = inner_product_real(rhk[1], rMiRk);
      bb = num/den;
      for (int j = 0; j < n2; j++) reDk[j] = rhk[0][j] + bb*reDk[j];
      cg_error = sqrt(inner_product_real(reRk,reRk));
      counter++;
      if ((counter%5)==0) printf("NV = %e CG Error = %e at %d step\n", NV, cg_error, counter);
      if (counter == 1000) break;

    }                                  // end of CG cycle
    if (counter == 1000) break;
    for (int j = 0; j < n2; j++) wrk->YYk[j] = wrk->YYk[j] + reDyk[j];
    printf("Error = %e at %3d CG step, H/L = %.14e Speed = %.12e\n", NV, counter, hmax, 1./sqrt(q));
    counter2++;
    writer_out(wrk->YYk,"stokes_spec.dat");
  }  // end of outer cycle

  if (counter == 1400) err_msg("CG diverged");

  ifft_even(wrk->YYk,wrk->YY);  
  writer_out(wrk->YY,"stokes_phys.dat");
  writer_out(wrk->YY,"restart.dat");

  FILE *fid = fopen("steep.data","a");
  fprintf(fid, "%.12e\t%.12e\n", hmax, 1./sqrt(q));
  fclose(fid);
  
}

/*void iterate(work_ptr wrk) {
  err = 1;
  int nn = 0;
  char st[80];

  write_out(wrk->Yk,"before.txt");
  while (err > tol) { 
    prepare_rhs(wrk); num = 0.; den = 0.;
    for (int j = 0; j < N; j++) {
      num += conj(wrk->Yk[j])*(hk[3][j]);
      den += mk[j]*pow(cabs(wrk->Yk[j]),2);
    }
    sp = 1.*num/den;
    err = 0.; l2 = 0; err2;
    for (int j = 0; j < N; j++) {
      rhs_yk = ((hk[3][j])/mk[j] - gama*sp*(wrk->Yk[j]) )*dtau;
      wrk->Yk[j] += rhs_yk;
      if (j > N/2) wrk->Yk[j] = conj(wrk->Yk[N-j]);
    }
    for (int j = 0; j < N; j++) {
      l2  += cpow(cabs(wrk->Yk[j]),2);
      err += cpow(cabs(rhs_yk),2);
      err2 += cpow(cabs(hk[3][j]),2);
    }
    err = sqrt(err/l2);
    err2 = sqrt(err2); 
    nn++;
    if ( (nn % 10)==0 ) {  
      printf("Iteration %5d. distance:  %e\tHeight:  %e\tSpeed: %.15e\tFluid Level: %e\n",nn, err,hmax/(2*pi), 1/sqrt(q),ml);
    }
  }
  printf("Iteration %5d. distance:  %e\tHeight:  %e\tSpeed: %.15e\tFluid Level: %e\n",nn, err,hmax/(2*pi), 1/sqrt(q),ml);
  ifft(wrk->Yk,wrk->Y);
  write_out(wrk->Y,"stokes_phys.dat");
  write_out(wrk->Y,"restart.dat");
  write_out(wrk->Yk,"stokes_spec.dat");

}*/





















